use crate::mq;
use axum::extract::ws::{Message, WebSocket};
use bytes::Bytes;
use futures_util::stream::{SplitSink, SplitStream};
use futures_util::{SinkExt, StreamExt};
use std::sync::Arc;
use tokio::sync::Mutex;
use tokio::sync::mpsc::{Receiver, Sender};

pub struct WsHandler {
    topic: String,
    mq_client_factory: Arc<mq::client::Factory>,
}

impl WsHandler {
    pub fn new(topic: String, mq_client_factory: Arc<mq::client::Factory>) -> Self {
        Self {
            topic,
            mq_client_factory,
        }
    }
    pub async fn handle(&self, socket: WebSocket) {
        let span = tracing::error_span!("WsHandler::handle", topic = self.topic.clone());
        let _guard = span.enter();

        let (sender, receiver) = tokio::sync::mpsc::channel(1024);

        let (mq_client, mut event_loop_task) = match self.mq_client_factory.create() {
            Ok(v) => v,
            Err(e) => {
                log::error!("创建新的mq客户端失败：{:?}", e);
                return;
            }
        };

        if let Err(e) = mq_client
            .subscribe(self.topic.clone(), Arc::new(MqPacketHandler::new(sender)))
            .await
        {
            log::error!("订阅失败{}：{:?}", self.topic, e);
            return;
        }

        let (ws_sender, ws_receiver) = socket.split();
        let ws_sender = Arc::new(Mutex::new(ws_sender));

        let mut ws_watch_handler =
            WsWatchHandler::new(self.topic.clone(), ws_receiver, ws_sender.clone());
        let mut mq_watch_handler =
            MqWatchHandler::new(self.topic.clone(), receiver, ws_sender.clone());

        let mq_event_loop_handle = tokio::spawn(async move {
            event_loop_task.run().await;
        });

        let mq_watch_handle = tokio::spawn(async move {
            mq_watch_handler.handle().await;
        });

        let ws_watch_handle = tokio::spawn(async move {
            ws_watch_handler.handle().await;
        });

        tokio::select! {
            _ = mq_event_loop_handle => {
                log::info!("mq事件循环任务结束");
            },
            _ = mq_watch_handle => {
                log::info!("mq监听任务结束");
            },
            _ = ws_watch_handle => {
                log::info!("websocket监听任务结束");
            },
        }

        if let Err(e) = mq_client.unsubscribe(self.topic.clone()).await {
            log::error!("取消订阅失败{}: {:?}", self.topic, e);
        }
    }
}

struct WsWatchHandler {
    topic: String,
    ws_receiver: SplitStream<WebSocket>,
    ws_sender: Arc<Mutex<SplitSink<WebSocket, Message>>>,
}

impl WsWatchHandler {
    fn new(
        topic: String,
        ws_receiver: SplitStream<WebSocket>,
        ws_sender: Arc<Mutex<SplitSink<WebSocket, Message>>>,
    ) -> Self {
        Self {
            topic,
            ws_receiver,
            ws_sender,
        }
    }
    async fn handle(&mut self) {
        let span = tracing::error_span!("WsWatchHandler", topic = self.topic);
        let _guard = span.enter();
        while let Some(result) = self.ws_receiver.next().await {
            match result {
                Ok(Message::Ping(_)) => {
                    log::debug!("接收到ping");
                    if let Err(e) = self
                        .ws_sender
                        .lock()
                        .await
                        .send(Message::Pong(Bytes::from("pong")))
                        .await
                    {
                        log::error!("发送pong失败：{:?}", e);
                        return;
                    }
                }
                Ok(Message::Pong(_)) => {
                    log::debug!("接收到pong");
                }
                Err(e) => {
                    log::error!("接收消息失败：{:?}", e);
                    return;
                }
                _ => {
                    log::debug!("其他消息：{:?}", result);
                }
            }
        }
    }
}

struct MqWatchHandler {
    topic: String,
    receiver: Receiver<Vec<u8>>,
    ws_sender: Arc<Mutex<SplitSink<WebSocket, Message>>>,
}

impl MqWatchHandler {
    fn new(
        topic: String,
        receiver: Receiver<Vec<u8>>,
        ws_sender: Arc<Mutex<SplitSink<WebSocket, Message>>>,
    ) -> Self {
        MqWatchHandler {
            topic,
            receiver,
            ws_sender,
        }
    }
    async fn send(&self, payload: Vec<u8>) -> Result<(), anyhow::Error> {
        let utf8bytes =
            axum::extract::ws::Utf8Bytes::try_from(payload).map_err(|e| anyhow::Error::new(e))?;

        let msg = Message::Text(utf8bytes);
        let mut websocket = self.ws_sender.lock().await;
        websocket
            .send(msg)
            .await
            .map_err(|e| anyhow::Error::new(e))?;
        Ok(())
    }

    async fn handle(&mut self) {
        let span = tracing::error_span!("MqWatchHandler", topic = self.topic);
        let _guard = span.enter();
        while let Some(data) = self.receiver.recv().await {
            if let Err(e) = self.send(data).await {
                log::error!("发送websocket失败： {:?}", e);
                break;
            }
        }
    }
}

pub struct MqPacketHandler {
    sender: Sender<Vec<u8>>,
}

impl MqPacketHandler {
    fn new(sender: Sender<Vec<u8>>) -> Self {
        MqPacketHandler { sender }
    }
}

#[async_trait::async_trait]
impl mq::client::event_loop_task::PacketHandler for MqPacketHandler {
    async fn handle(&self, packet: rumqttc::Publish) {
        if let Err(e) = self.sender.send(packet.payload.to_vec()).await {
            log::error!(
                "数据写入channel失败, topic={}, 错误： {:?}",
                packet.topic,
                e
            );
        }
    }
}
